#include <windows.h>
#include <shlobj.h>
#include <string>
#include <vector>
#include <algorithm>
#include <objbase.h>
#include <shlwapi.h>
#include <comdef.h>
#include <sstream>
#define CLASS_E_CLASSNOTFOUND _HRESULT_TYPEDEF_(0x80040111L)
#define IDM_DISPLAY_PATHS 0
#define IDM_GET_PDF_COUNT 1 // 新增菜单ID
#define OUTPUTWSTREAM(wss)  OutputDebugStringW(wss.str().c_str())

WINBOOL runAndWait(const std::wstring cmdLine, bool hideWindow = false) {
    STARTUPINFOW si = {};
    PROCESS_INFORMATION pi = {};
    si.cb = sizeof(si);

    if (hideWindow) {
        si.dwFlags = STARTF_USESHOWWINDOW;
        si.wShowWindow = SW_HIDE;
    }

    auto res = CreateProcessW(NULL, (LPWSTR)cmdLine.c_str(), NULL, NULL, FALSE, 0, NULL, NULL,
                   &si, &pi);
    if (res){
        WaitForSingleObject(pi.hProcess, INFINITE);
        CloseHandle(pi.hProcess);
        CloseHandle(pi.hThread);
    }
    else
        return false;

    return true;
}

const std::wstring get_temp_file_name_w()
{
    WCHAR tempPath[MAX_PATH];
    WCHAR tempFileName[MAX_PATH];

    // 获取系统临时目录
    GetTempPathW(MAX_PATH, tempPath);

    // 生成临时文件名
    UINT result = GetTempFileNameW(
        tempPath,    // 目录路径
        L"YHDN",     // 前缀
        0,           // 0表示使用随机数
        tempFileName // 输出缓冲区
    );

    return tempFileName;
}

// 全局变量
HINSTANCE g_hInst = NULL;
const CLSID CLSID_PdfContextMenu = {0x1D4988A1, 0xB9A5, 0x4E5D, {0x8A, 0x3C, 0x3A, 0x5E, 0x9D, 0x1F, 0x3B, 0x7A}};

class CPdfContextMenu : public IShellExtInit, public IContextMenu
{
private:
    LONG m_cRef;
    std::vector<std::wstring> m_files;

public:
    CPdfContextMenu() : m_cRef(1) {}

    // IUnknown methods
    STDMETHODIMP QueryInterface(REFIID riid, void **ppv)
    {
        if (riid == IID_IUnknown || riid == IID_IShellExtInit)
        {
            *ppv = static_cast<IShellExtInit *>(this);
        }
        else if (riid == IID_IContextMenu)
        {
            *ppv = static_cast<IContextMenu *>(this);
        }
        else
        {
            *ppv = NULL;
            return E_NOINTERFACE;
        }
        AddRef();
        return S_OK;
    }

    STDMETHODIMP_(ULONG)
    AddRef()
    {
        return InterlockedIncrement(&m_cRef);
    }

    STDMETHODIMP_(ULONG)
    Release()
    {
        ULONG cRef = InterlockedDecrement(&m_cRef);
        if (cRef == 0)
        {
            delete this;
        }
        return cRef;
    }

    // IShellExtInit method
    STDMETHODIMP Initialize(LPCITEMIDLIST pidlFolder, LPDATAOBJECT pDataObj, HKEY hkeyProgID)
    {
        if (pDataObj == NULL)
            return E_INVALIDARG;

        FORMATETC fmt = {CF_HDROP, NULL, DVASPECT_CONTENT, -1, TYMED_HGLOBAL};
        STGMEDIUM stg = {TYMED_HGLOBAL};
        HDROP hDrop;

        if (FAILED(pDataObj->GetData(&fmt, &stg)))
            return E_FAIL;

        hDrop = (HDROP)GlobalLock(stg.hGlobal);
        if (hDrop == NULL)
            return E_FAIL;

        // 获取文件数量
        UINT nFiles = DragQueryFileW(hDrop, 0xFFFFFFFF, NULL, 0);
        m_files.clear();

        // 检查所有文件是否都是PDF
        bool allPdf = true;
        for (UINT i = 0; i < nFiles; i++)
        {
            WCHAR szFile[MAX_PATH];
            DragQueryFileW(hDrop, i, szFile, MAX_PATH);

            // 检查文件扩展名
            std::wstring file(szFile);
            size_t dotPos = file.find_last_of(L".");
            if (dotPos == std::wstring::npos ||
                _wcsicmp(file.substr(dotPos).c_str(), L".pdf") != 0)
            {
                allPdf = false;
                break;
            }

            m_files.push_back(file);
        }

        GlobalUnlock(stg.hGlobal);
        ReleaseStgMedium(&stg);

        // 如果不是所有文件都是PDF，清空列表
        if (!allPdf)
        {
            m_files.clear();
        }

        return allPdf ? S_OK : E_FAIL;
    }

    // IContextMenu methods
    STDMETHODIMP QueryContextMenu(HMENU hmenu, UINT indexMenu, UINT idCmdFirst, UINT idCmdLast, UINT uFlags)
    {
        if ((uFlags & CMF_DEFAULTONLY) || m_files.empty())
            return MAKE_HRESULT(SEVERITY_SUCCESS, 0, 0);

        // 添加第一个菜单项：合并PDF文件
        InsertMenuW(hmenu, indexMenu, MF_BYPOSITION | MF_STRING,
                    idCmdFirst + IDM_DISPLAY_PATHS, L"票据PDF打印");

        // // 添加第二个菜单项：获取PDF文件个数
        // std::wstring countMenuText = L"合并打印所选的PDF文件";
        // InsertMenuW(hmenu, indexMenu + 1, MF_BYPOSITION | MF_STRING,
        //             idCmdFirst + IDM_GET_PDF_COUNT, countMenuText.c_str());

        return MAKE_HRESULT(SEVERITY_SUCCESS, 0, 1); // 返回2个菜单项
    }

    STDMETHODIMP InvokeCommand(LPCMINVOKECOMMANDINFO pici)
    {
        if (HIWORD(pici->lpVerb) != 0)
            return E_FAIL;

        UINT idCmd = LOWORD(pici->lpVerb);

        if (idCmd == IDM_DISPLAY_PATHS)
        {
            // 第一个菜单项的功能：合并PDF文件
            std::wstringstream wstream;
            wstream << L"pdftk.exe";
            for(auto filePath : m_files){

                wstream << " \"" << filePath << "\"";
            }
            auto temp_file_name = get_temp_file_name_w() + L".pdf";
            wstream << " cat output " << temp_file_name;
            OutputDebugStringW(wstream.str().c_str());
            //执行命令
            auto res = runAndWait(wstream.str(),true);
            //用adrord32.exe打开
            std::wstringstream wstream2;
            wstream2 << "C:\\Program Files (x86)\\Adobe\\Reader 9.0\\Reader\\AcroRd32.exe " << temp_file_name;
            OutputDebugStringW(wstream2.str().c_str());            
            res = runAndWait(wstream2.str(),true);
            
        }
        else if (idCmd == IDM_GET_PDF_COUNT)
        {
            // 第一个菜单项的功能：合并PDF文件
            // std::wstringstream wstream;
            // wstream << L"pdftk.exe";
            // for(auto filePath : m_files){

            //     wstream << " \"" << filePath << "\"";
            // }
            // auto temp_file_name = get_temp_file_name_w();
            // auto temp_file_name2 = temp_file_name + L".pdf";
            // wstream << " cat output " << temp_file_name;
            // OutputDebugStringW(wstream.str().c_str());
            // //执行命令
            // auto res = runAndWait(wstream.str(),true);

            // //N-UP
            // std::wstringstream wstream2;
            // wstream2 << L"nup_pdf.exe \"" << temp_file_name << L"\" \"" << temp_file_name2 << L"\"" << " 2 -k 0";
            // OutputDebugStringW(wstream2.str().c_str());
            // res = runAndWait(wstream2.str(),true);
            // //Sumatrapdf打印 
            // std::wstringstream wstream3;
            // wstream3 << L"sumatrapdf.exe -print-to-default \"" << temp_file_name2 << L"\"";
            // OutputDebugStringW(wstream3.str().c_str());
            // res = runAndWait(wstream3.str(),true);
            OutputDebugStringA("合并并打印");



        }
        else
        {
            return E_FAIL;
        }

        return S_OK;
    }

    STDMETHODIMP GetCommandString(UINT_PTR idCmd, UINT uType, UINT *pReserved, LPSTR pszName, UINT cchMax)
    {
        if (idCmd != IDM_DISPLAY_PATHS && idCmd != IDM_GET_PDF_COUNT)
            return E_INVALIDARG;

        if (uType == GCS_HELPTEXT)
        {
            if (pszName)
            {
                std::wstring helpText;
                if (idCmd == IDM_DISPLAY_PATHS)
                {
                    helpText = L"合并选中的PDF文件并打开";
                }
                else if (idCmd == IDM_GET_PDF_COUNT)
                {
                    helpText = L"打印PDF开票";
                }
                lstrcpynW((LPWSTR)pszName, helpText.c_str(), cchMax);
            }
            return S_OK;
        }

        return E_NOTIMPL;
    }
};

// Class factory
class CPdfContextMenuClassFactory : public IClassFactory
{
public:
    // IUnknown
    STDMETHODIMP QueryInterface(REFIID riid, void **ppv)
    {
        if (riid == IID_IUnknown || riid == IID_IClassFactory)
        {
            *ppv = static_cast<IClassFactory *>(this);
            return S_OK;
        }
        *ppv = NULL;
        return E_NOINTERFACE;
    }

    STDMETHODIMP_(ULONG)
    AddRef()
    {
        return 2;
    }

    STDMETHODIMP_(ULONG)
    Release()
    {
        return 1;
    }

    // IClassFactory
    STDMETHODIMP CreateInstance(IUnknown *pUnkOuter, REFIID riid, void **ppv)
    {
        if (pUnkOuter != NULL)
            return CLASS_E_NOAGGREGATION;

        CPdfContextMenu *pPdfContextMenu = new (std::nothrow) CPdfContextMenu();
        if (pPdfContextMenu == NULL)
            return E_OUTOFMEMORY;

        HRESULT hr = pPdfContextMenu->QueryInterface(riid, ppv);
        pPdfContextMenu->Release();
        return hr;
    }

    STDMETHODIMP LockServer(BOOL fLock)
    {
        return S_OK;
    }
};

// DLL exports
STDAPI DllGetClassObject(REFCLSID rclsid, REFIID riid, void **ppv)
{
    if (!IsEqualCLSID(rclsid, CLSID_PdfContextMenu))
        return CLASS_E_CLASSNOTFOUND;

    CPdfContextMenuClassFactory *pcf = new (std::nothrow) CPdfContextMenuClassFactory();
    if (pcf == NULL)
        return E_OUTOFMEMORY;

    HRESULT hr = pcf->QueryInterface(riid, ppv);
    pcf->Release();
    return hr;
}

STDAPI DllCanUnloadNow()
{
    return S_OK;
}

STDAPI DllRegisterServer()
{
    WCHAR szModule[MAX_PATH];
    GetModuleFileNameW(g_hInst, szModule, MAX_PATH);

    // 注册CLSID
    WCHAR szCLSID[64];
    StringFromGUID2(CLSID_PdfContextMenu, szCLSID, ARRAYSIZE(szCLSID));

    WCHAR szKey[256];
    HKEY hKey;

    // 创建CLSID键
    wsprintfW(szKey, L"CLSID\\%s", szCLSID);
    if (RegCreateKeyW(HKEY_CLASSES_ROOT, szKey, &hKey) != ERROR_SUCCESS)
        return SELFREG_E_CLASS;
    RegSetValueExW(hKey, NULL, 0, REG_SZ, (BYTE *)L"PDF Context Menu", sizeof(L"PDF Context Menu"));
    RegCloseKey(hKey);

    // 注册InprocServer32
    wsprintfW(szKey, L"CLSID\\%s\\InprocServer32", szCLSID);
    if (RegCreateKeyW(HKEY_CLASSES_ROOT, szKey, &hKey) != ERROR_SUCCESS)
        return SELFREG_E_CLASS;
    RegSetValueExW(hKey, NULL, 0, REG_SZ, (BYTE *)szModule, (lstrlenW(szModule) + 1) * sizeof(WCHAR));
    RegSetValueExW(hKey, L"ThreadingModel", 0, REG_SZ, (BYTE *)L"Apartment", sizeof(L"Apartment"));
    RegCloseKey(hKey);

    // 关联到所有文件类型
    if (RegCreateKeyW(HKEY_CLASSES_ROOT, L"*\\shellex\\ContextMenuHandlers\\PDFContextMenu", &hKey) != ERROR_SUCCESS)
        return SELFREG_E_CLASS;
    RegSetValueExW(hKey, NULL, 0, REG_SZ, (BYTE *)szCLSID, (lstrlenW(szCLSID) + 1) * sizeof(WCHAR));
    RegCloseKey(hKey);

    return S_OK;
}

STDAPI DllUnregisterServer()
{
    WCHAR szCLSID[64];
    StringFromGUID2(CLSID_PdfContextMenu, szCLSID, ARRAYSIZE(szCLSID));

    WCHAR szKey[256];
    std::wstringstream wss;

    // 删除上下文菜单处理器注册
    RegDeleteKeyW(HKEY_CLASSES_ROOT, L"*\\shellex\\ContextMenuHandlers\\PDFContextMenu");

    // 删除CLSID键
    wss << L"CLSID\\" << szCLSID <<"\\InprocServer32";
    // wsprintfW(szKey, L"CLSID\\%s\\InprocServer32", szCLSID);
    RegDeleteKeyW(HKEY_CLASSES_ROOT, wss.str().c_str());
    OutputDebugStringW(L"DllUnregisterServer");
    OUTPUTWSTREAM(wss);
    
    
    wss.str(L"");
    wss.clear();
    
    wss << L"CLSID\\" << szCLSID;
    RegDeleteKeyW(HKEY_CLASSES_ROOT, wss.str().c_str());
    OutputDebugStringW(L"DllUnregisterServer");
    OUTPUTWSTREAM(wss);

    return S_OK;
}

BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved)
{
    switch (ul_reason_for_call)
    {
    case DLL_PROCESS_ATTACH:
    {
        g_hInst = hModule;
        DisableThreadLibraryCalls(hModule);
        // 获取当前模块路径
        WCHAR s_buffer[MAX_PATH] = {0};
        GetModuleFileNameW(hModule, s_buffer, MAX_PATH);
        std::wstringstream wss;
        wss << "dllmain函数 当前模块路径:" << s_buffer;
        OutputDebugStringW(s_buffer);

        
    }
    break;
    case DLL_THREAD_ATTACH:
    case DLL_THREAD_DETACH:
    case DLL_PROCESS_DETACH:
        break;
    }
    return TRUE;
}